import re
from typing import TypedDict, Annotated
from langgraph.graph import StateGraph
from nodes import NovelNodes

class AgentState(TypedDict):
    ref: Annotated[str, "参考资料"]
    outline: Annotated[str, "小说大纲内容"]
    refine_count: Annotated[int, "大纲迭代次数"]
    current_chapter: Annotated[int, "当前撰写章节编号（从1开始）"]
    chapter_content: Annotated[str, "章节内容"]
    word_count: Annotated[int, "章节字数"]

class NovelWorkflow:
    def __init__(self):
        self.nodes = NovelNodes()
        self.graph = StateGraph(AgentState)

    def _should_continue_writing(self, state: AgentState) -> str:
        total_chapters = re.search(r'总章节数：(\d+)章', state['outline']).group(1)
        total_chapters = int(total_chapters)

        print(f"当前章节: {state['current_chapter']}, 总章节数: {total_chapters}")  # Debug print
        print('------------------------------')
        return "continue" if state["current_chapter"] <= total_chapters else "review"

    def _should_refine_outline(self, state: AgentState) -> str:
        return "refine" if state["refine_count"] < 2 else "write"

    def compile_workflow(self):
        # 添加工作流节点
        self.graph.add_node("generate_outline", self.nodes.generate_initial_outline)
        self.graph.add_node("refine_outline", self.nodes.refine_outline)
        self.graph.add_node("write_chapter", self.nodes.write_chapter)
        self.graph.add_node("final_review", self.nodes.final_review)

        # 设置初始流程
        self.graph.set_entry_point("generate_outline")

        # 大纲迭代条件分支
        self.graph.add_conditional_edges(
            "generate_outline",
            lambda s: "refine" if s["refine_count"] == 0 else "write",
            {
                "refine": "refine_outline",
                "write": "write_chapter"
            }
        )

        # 构建完整流程
        self.graph.add_edge("refine_outline", "write_chapter")
        self.graph.add_conditional_edges(
            "write_chapter",
            self._should_continue_writing,
            {
                "continue": "write_chapter",
                "review": "final_review"
            }
        )
        self.graph.set_finish_point("final_review")

        # 设置递归限制
        return self.graph.compile(checkpointer=None)